#include <cuda_runtime.h>
#include "freshman.h"
#include <cstdio>

void sumMatrixCPU(float *matA, float *matB, float *matC, int nx, int ny) {
    float *a = matA;
    float *b = matB;
    float *c = matC;

    for (int j = 0; j < ny; j++) {
        for (int i = 0; i < nx; i++) {
            c[i] = a[i] + b[i];
        }
        a += nx;
        b += nx;
        c += nx;
    }
}


__global__ void sumMatrix(const float *matA, const float *matB, float *matC, int nx, int ny) {
    int ix = blockIdx.x * blockDim.x + threadIdx.x;
    int iy = blockIdx.y * blockDim.y + threadIdx.y;

    int idx = ix + iy * ny;

    if (ix < nx && iy < ny) {
        matC[idx] = matA[idx] + matB[idx];
    }
}

int main(int argc, char **argv) {
    printf("starting....\n");

    initDevice(0);

    int nx = 1 << 14;
    int ny = 1 << 14;
    int nxy = nx * ny;
    int nBytes = nxy * sizeof(float);

    // Malloc
    float *A_host = (float *) malloc(nBytes);
    float *B_host = (float *) malloc(nBytes);
    float *C_host = (float *) malloc(nBytes);
    float *C_from_GPU = (float *) malloc(nBytes);

    // CudaMalloc
    float *A_dev = nullptr;
    float *B_dev = nullptr;
    float *C_dev = nullptr;
    CHECK(cudaMalloc((void **) &A_dev, nBytes));
    CHECK(cudaMalloc((void **) &B_dev, nBytes));
    CHECK(cudaMalloc((void **) &C_dev, nBytes));

    // CudaMemcpy
    CHECK(cudaMemcpy(A_dev, A_host, nBytes, cudaMemcpyHostToDevice));
    CHECK(cudaMemcpy(B_dev, B_host, nBytes, cudaMemcpyHostToDevice));

    int dimx = 8;
    int dimy = 8;

    // cpu compute
//    CHECK(cudaMemcpy(C_from_GPU, C_dev, nBytes, cudaMemcpyDeviceToDevice));
    double iStart = cpuSecond();
    sumMatrixCPU(A_host, B_host, C_host, nx, ny);
    double iElapse = cpuSecond() - iStart;
    printf("CPU Execution Time elapsed %f sec\n", iElapse);

    // 2d block and 2d grid
    dim3 block_0(dimx, dimy);
    dim3 grid_0((nx - 1) / block_0.x + 1, (ny - 1) / block_0.y + 1);
    iStart = cpuSecond();
    sumMatrix<<<grid_0, block_0>>>(A_dev, B_dev, C_dev, nx, ny);
    CHECK(cudaDeviceSynchronize());
    iElapse = cpuSecond() - iStart;
    printf("GPU Execution configuration<<<(%d,%d),(%d,%d)>>> Time elapsed %f sec\n",
           grid_0.x, grid_0.y, block_0.x, block_0.y, iElapse);
    CHECK(cudaMemcpy(C_from_GPU, C_dev, nBytes, cudaMemcpyDeviceToHost));
    checkResult(C_host, C_from_GPU, nxy);

    cudaFree(A_dev);
    cudaFree(B_dev);
    cudaFree(C_dev);
    free(A_host);
    free(B_host);
    free(C_host);
    free(C_from_GPU);
    cudaDeviceReset();
    return 0;
}
